import torch
import torch.nn as nn

def sparse_max_pool(input, size):
    positive = (input > 0).float()
    negative = (input < 0).float()
    output = nn.functional.adaptive_max_pool2d(input * positive, size) - nn.functional.adaptive_max_pool2d(-input * negative, size)
    return output


def multiscaleEPE(network_output, target_image, weights=None, L_model=None):
    def one_scale(output, target, L_model):
        b, chn, h, w = output.size()

        target_scaled = nn.functional.adaptive_avg_pool3d(target, (chn, h, w))
        # print(target_scaled.size())

        return L_model(output, target_scaled) # EPE(output, target_scaled, L_model)
    

    if type(network_output) not in [tuple, list]:
        network_output = [network_output]
    if weights is None:
        weights = [ 1.0/32,1.0/16.0, 1.0/8.0, 1.0/4.0, 1.0/2.0]  
    assert(len(weights) == len(network_output))

    loss = 0
    basic_loss = 0
    color_loss = 0
    edge_loss = 0
    for output, weight in zip(network_output, weights):
        lossinfo = one_scale(output, target_image, L_model)
        loss += weight * lossinfo[0]
        basic_loss += weight * lossinfo[1]
        color_loss += weight * lossinfo[2]
        edge_loss += weight * lossinfo[3]
        edge_residue = lossinfo[4].cpu()
        
    
    return loss, basic_loss, color_loss, edge_loss, edge_residue


